#!/usr/bin/python3
import os
import cv2
import numpy as np
import pandas as pd

face_attr_list_file = r'/rootfs/media/kasim/Data1/data/VideoFaceImage/extract_image_list_72.attr.txt'

# VERY_GOOD = 0
# GOOD = 1
# BAD = 2
# VERY_BAD = 3
filter_quality = 3
filter_quality_thr = 0.25

min_pitch = -18
max_pitch = 18
min_roll = -18
max_roll = 18
min_yaw = -18
max_yaw = 18

min_angle_score = 0.1
max_angle_diff = 20


def main():
    # file_name;bin_pitch;bin_roll;bin_yaw;pitch;roll;yaw;quality;mask
    attr_cvs = pd.read_csv(face_attr_list_file, sep=';')
    file_names = attr_cvs['file_name']
    qualitys = attr_cvs['quality']
    bin_pitchs = attr_cvs['bin_pitch']
    bin_rolls = attr_cvs['bin_roll']
    bin_yaws = attr_cvs['bin_yaw']
    pitchs = attr_cvs['pitch']
    rolls = attr_cvs['roll']
    yaws = attr_cvs['yaw']

    out_file_path = os.path.join(os.path.dirname(face_attr_list_file), 'image_list.attr.txt')           # 分类角度与回归角度差别较大的，或分类角度置信度较低的
    out_file1_path = os.path.join(os.path.dirname(face_attr_list_file), 'image_list_diff.attr.txt')     # 分类角度与回归角度差别较大的
    out_file2_path = os.path.join(os.path.dirname(face_attr_list_file), 'image_list_score.attr.txt')    # 分类角度置信度较低的
    out_file3_path = os.path.join(os.path.dirname(face_attr_list_file), 'image_list_both.attr.txt')     # 分类角度与回归角度差别较大的，且分类角度置信度较低的
    out_file = open(out_file_path, 'w')
    out_file1 = open(out_file1_path, 'w')
    out_file2 = open(out_file2_path, 'w')
    out_file3 = open(out_file3_path, 'w')
    out_file.write('file_name;bin_pitch;bin_roll;bin_yaw;pitch;roll;yaw;quality')
    out_file1.write('file_name;bin_pitch;bin_roll;bin_yaw;pitch;roll;yaw;quality')
    out_file2.write('file_name;bin_pitch;bin_roll;bin_yaw;pitch;roll;yaw;quality')
    out_file3.write('file_name;bin_pitch;bin_roll;bin_yaw;pitch;roll;yaw;quality')
    out_count = 0
    out_count1 = 0
    out_count2 = 0
    out_count3 = 0
    file_count = 0
    for file_name, quality_s, bin_pitch_s, bin_roll_s, bin_yaw_s, pitch_s, roll_s, yaw_s in zip(file_names, qualitys, bin_pitchs, bin_rolls, bin_yaws, pitchs, rolls, yaws):
        file_count += 1
        if file_count % 10000 == 0:
            print('Process Filter File Count: {}'.format(file_count))
        quality = list(map(float, quality_s.split(',')))
        quality = np.array(quality, dtype=np.float32)
        max_idx = np.argmax(quality)
        if (max_idx == filter_quality) and (quality[max_idx] >= filter_quality_thr):
            continue

        pitch = float(pitch_s)
        roll = float(roll_s)
        yaw = float(yaw_s)

        bin_pitch = bin_pitch_s.split(',')
        bin_roll = bin_roll_s.split(',')
        bin_yaw = bin_yaw_s.split(',')

        bin_pitch_score = float(bin_pitch[1])
        bin_pitch = float(bin_pitch[0])
        bin_roll_score = float(bin_roll[1])
        bin_roll = float(bin_roll[0])
        bin_yaw_score = float(bin_yaw[1])
        bin_yaw = float(bin_yaw[0])

        pitch_diff = abs(bin_pitch - pitch)
        roll_diff = abs(bin_roll - roll)
        yaw_diff = abs(bin_yaw - yaw)
        has_write = False
        if max_angle_diff <= pitch_diff or max_angle_diff <= roll_diff or max_angle_diff <= yaw_diff:
            out_info = '{};{};{};{};{};{};{};{}\n'.format(file_name, bin_pitch_s, bin_roll_s, bin_yaw_s, pitch_s, roll_s, yaw_s, quality_s)
            out_file.write(out_info)
            out_count += 1
            has_write = True
            out_file1.write(out_info)
            out_count1 += 1

        if pitch < min_pitch or max_pitch < pitch or roll < min_roll or max_roll < roll or yaw < min_yaw or max_yaw < yaw:
            continue

        if bin_pitch_score <= min_angle_score or bin_roll_score <= min_angle_score or bin_yaw_score <= min_angle_score:
            out_info = '{};{};{};{};{};{};{};{}\n'.format(file_name, bin_pitch_s, bin_roll_s, bin_yaw_s, pitch_s, roll_s, yaw_s, quality_s)
            if not has_write:
                out_file.write(out_info)
                out_count += 1
                has_write = True
            else:
                out_file3.write(out_info)
                out_count3 += 1
            out_file2.write(out_info)
            out_count2 += 1
    print('Process Filter File Count: {}'.format(file_count))
    out_file.close()
    out_file1.close()
    out_file2.close()
    out_file3.close()
    os.system('chmod a+wr {}'.format(out_file_path))
    os.system('chmod a+wr {}'.format(out_file1_path))
    os.system('chmod a+wr {}'.format(out_file2_path))
    os.system('chmod a+wr {}'.format(out_file3_path))
    print('out_count:', out_count, 'out_count1:', out_count1, 'out_count2:', out_count2, 'out_count3:', out_count3)


if __name__ == '__main__':
    main()
